import json

from app.provider.imgsearch_provider import ImgsearchProvider
from zhanshop.hash import Hash
from zhanshop.json import Json


class ImageService():
    models = None

    def __init__(self):
        self.models = {}

    def getModel(self, database, table, modelParam):
        """
        获取模型
        :param database:
        :param table:
        :param modelParam:
        :return:
        """
        key = Hash.md5(Json.encode(modelParam))
        model = self.models.get(key, None)
        if (model == None):
            model = ImgsearchProvider(database, table, modelParam)
            self.models[key] = model
        return model

    def getRepository(self, database: str, table: str, imgPath: str, limit: int = 100, modelParam: json = None):
        """
        搜索相似图像
        :param database:
        :param table:
        :param documents:
        :param modelParam:
        :return:
        """
        if (modelParam == None or modelParam == ""): modelParam = {
            "model": "all-MiniLM-L6-v2",
            "chromadb": {"metadata": {"hnsw:space": "cosine"}}
        }
        model = self.getModel(database, table, modelParam)
        return model.query(imgPath, limit)

    def postRepository(self, database: str, table: str, imgPaths: list, metadatas: list, modelParam: json = None):
        """
        将图片添加到图像库
        :param database:
        :param table:
        :param imgPaths:
        :param metadatas:
        :param modelParam:
        :return:
        """
        if (modelParam == None or modelParam == ""): modelParam = {
            "model": "all-MiniLM-L6-v2",
            "chromadb": {"metadata": {"hnsw:space": "cosine"}}
        }
        model = self.getModel(database, table, modelParam)
        model.batchAdd(imgPaths, metadatas)
        return {}